"""
# @Time    : 2021/7/2 5:22 下午
# @Author  : hezhiqiang
# @Email   : tinyzqh@163.com
# @File    : env_discrete.py
"""

import gym
from gym import spaces
import numpy as np
from envs.env_core import EnvCore


class DiscreteActionEnv(object):
    """
    对于离散动作环境的封装
    """

    def __init__(self):
        """
        初始化离散动作环境
        """
        self.env = EnvCore()  # 创建核心环境
        self.num_agent = self.env.agent_num  # 智能体数量

        self.signal_obs_dim = self.env.obs_dim  # 单个智能体的观察维度
        self.signal_action_dim = self.env.action_dim  # 单个智能体的动作维度

        # 如果为真，动作是一个0...N的数字，否则动作是一个one-hot的N维向量
        self.discrete_action_input = False

        self.movable = True  # 是否可移动

        # 配置空间
        self.action_space = []  # 动作空间列表
        self.observation_space = []  # 观察空间列表
        self.share_observation_space = []  # 共享观察空间列表

        share_obs_dim = 0  # 共享观察维度
        total_action_space = []  # 总动作空间
        for agent_idx in range(self.num_agent):
            # 物理动作空间
            u_action_space = spaces.Discrete(self.signal_action_dim)  # 5个离散的动作

            # 如果可移动
            total_action_space.append(u_action_space)

            # 总动作空间
            # if len(total_action_space) > 1:
            #     # 所有动作空间都是离散的，所以简化为MultiDiscrete动作空间
            #     if all(
            #         [
            #             isinstance(act_space, spaces.Discrete)
            #             for act_space in total_action_space
            #         ]
            #     ):
            #         act_space = MultiDiscrete(
            #             [[0, act_space.n - 1] for act_space in total_action_space]
            #         )
            #     else:
            #         act_space = spaces.Tuple(total_action_space)
            # self.action_space.append(act_space)
            # else:
            self.action_space.append(total_action_space[agent_idx])  # 添加动作空间

            # 观察空间
            share_obs_dim += self.signal_obs_dim  # 累加共享观察维度
            self.observation_space.append(
                spaces.Box(
                    low=-np.inf,
                    high=+np.inf,
                    shape=(self.signal_obs_dim,),
                    dtype=np.float32,
                )
            )  # [-inf,inf]

        # 共享观察空间
        self.share_observation_space = [
            spaces.Box(low=-np.inf, high=+np.inf, shape=(share_obs_dim,), dtype=np.float32)
            for _ in range(self.num_agent)
        ]

    def step(self, actions):
        """
        执行环境步进
        
        输入actions维度假设：
        # actions shape = (5, 2, 5)
        # 5个线程的环境，里面有2个智能体，每个智能体的动作是一个one_hot的5维编码
        """

        results = self.env.step(actions)  # 执行环境步进
        obs, rews, dones, infos = results  # 解包结果
        return np.stack(obs), np.stack(rews), np.stack(dones), infos  # 返回结果

    def reset(self):
        """
        重置环境
        
        返回:
            初始观察
        """
        obs = self.env.reset()  # 重置环境
        return np.stack(obs)  # 返回堆叠的观察

    def close(self):
        """
        关闭环境
        """
        pass

    def render(self, mode="rgb_array"):
        """
        渲染环境
        
        参数:
            mode: 渲染模式
        """
        pass

    def seed(self, seed):
        """
        设置随机种子
        
        参数:
            seed: 随机种子
        """
        pass


class MultiDiscrete:
    """
    - 多离散动作空间由一系列具有不同参数的离散动作空间组成
    - 它可以适应离散动作空间或连续（Box）动作空间
    - 它对于表示游戏控制器或键盘很有用，其中每个键可以表示为一个离散动作空间
    - 它通过传递一个包含每个离散动作空间的[min, max]的数组数组来参数化
       其中离散动作空间可以取从`min`到`max`的任何整数（包括两者）
    注意：值0始终需要表示NOOP动作。
    例如：任天堂游戏控制器
    - 可以概念化为3个离散动作空间：
        1) 方向键：离散5 - NOOP[0], UP[1], RIGHT[2], DOWN[3], LEFT[4] - 参数：min: 0, max: 4
        2) A按钮：离散2 - NOOP[0], Pressed[1] - 参数：min: 0, max: 1
        3) B按钮：离散2 - NOOP[0], Pressed[1] - 参数：min: 0, max: 1
    - 可以初始化为
        MultiDiscrete([ [0,4], [0,1], [0,1] ])
    """

    def __init__(self, array_of_param_array):
        """
        初始化多离散动作空间
        
        参数:
            array_of_param_array: 参数数组的数组
        """
        super().__init__()
        self.low = np.array([x[0] for x in array_of_param_array])  # 最小值数组
        self.high = np.array([x[1] for x in array_of_param_array])  # 最大值数组
        self.num_discrete_space = self.low.shape[0]  # 离散空间数量
        self.n = np.sum(self.high) + 2  # 总动作数

    def sample(self):
        """
        返回一个数组，其中包含每个离散动作空间的一个样本
        """
        # 对于每一行：round(random .* (max - min) + min, 0)
        random_array = np.random.rand(self.num_discrete_space)  # 生成随机数组
        return [int(x) for x in np.floor(np.multiply((self.high - self.low + 1.0), random_array) + self.low)]  # 计算样本

    def contains(self, x):
        """
        检查x是否在动作空间中
        
        参数:
            x: 要检查的动作
            
        返回:
            是否在动作空间中
        """
        return (
            len(x) == self.num_discrete_space  # 长度相等
            and (np.array(x) >= self.low).all()  # 所有元素大于等于最小值
            and (np.array(x) <= self.high).all()  # 所有元素小于等于最大值
        )

    @property
    def shape(self):
        """
        返回动作空间的形状
        """
        return self.num_discrete_space

    def __repr__(self):
        """
        返回动作空间的字符串表示
        """
        return "MultiDiscrete" + str(self.num_discrete_space)

    def __eq__(self, other):
        """
        检查两个动作空间是否相等
        
        参数:
            other: 另一个动作空间
            
        返回:
            是否相等
        """
        return np.array_equal(self.low, other.low) and np.array_equal(self.high, other.high)


if __name__ == "__main__":
    """
    如果直接运行此文件，则执行步进
    """
    DiscreteActionEnv().step(actions=None)
